526B - Om Nom and Dark Park - CodeForces Solution

from typing import List, Union
from collections import namedtuple, deque
from datetime import datetime
import sys
import traceback
from time import perf_counter


class TreeNode:
    def __init__(self, val=0):
        self.val = val
        self.left = None
        self.right = None



class Solution:
    def __init__(self):
        self.count = 0

    def dark_park(self, n: int, lights: List) -> int:
        self.n = n
        self.lights = lights

        result = self.dfs(1)
        print(self.count)
        return self.count

    def dfs(self, node):
        if node >= len(self.lights) + 2:
            return 0

        index = node - 2
        if index == -1:
                                    v = 0
        else:
            v = self.lights[index]

        q = self.dfs(2 * node)
        r = self.dfs(2 * node + 1)
        v += max(q, r)
        self.count += abs(q - r)
        return v


TestCase = namedtuple('TestCase', 'n lights correct')


def read_test_cases(input_file, output_file):
    test_cases: List[TestCase] = []
    try:
        with open(input_file, 'r') as in_f:
            with open(output_file, 'r') as out_f:
                test_num = int(in_f.readline().strip())
                for _ in range(test_num):
                    n = int(in_f.readline().strip())
                    lights = in_f.readline().strip().split(' ')
                    lights = [int(i) for i in lights]
                    correct = int(out_f.readline().strip())
                    t = TestCase(n=n, lights=lights, correct=correct)
                    test_cases.append(t)
    except Exception as exc:
        exc_name = exc.__class__.__name__
        exc_msg = str(exc)
        exc_info = sys.exc_info()
        print('EXCEPTION:', exc_name, exc_msg)
        traceback.print_exception(*exc_info)
    return test_cases


def run_test_cases(test_cases: List[TestCase]):
    for t in test_cases:
        result = Solution().dark_park(n=t.n, lights=t.lights)
        print('N:', t.n, 'LIGHTS:', t.lights, 'CORRECT:', t.correct, 'RESULT:', result, 'CHECK:', result==t.correct)


if __name__ == '__main__':
    if '--debug' in sys.argv:
        start_time = datetime.utcnow().strftime('%Y.%m.%D %H:%M:%S')
        print('STARTING:', start_time)
        test_cases = read_test_cases('data/input.txt', 'data/output.txt')
        start_counter = perf_counter()
        run_test_cases(test_cases)
        stop_counter = perf_counter()
        print('COUNTER:', stop_counter - start_counter)
    else:
        n = int(input().strip())
        lights = input().strip().split(' ')
        lights = [int(i) for i in lights]
        Solution().dark_park(n=n, lights=lights)